#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/8/19
# @USER    : Shengji He
# @File    : dataset.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import torch
import torchvision
from torchvision import transforms as transforms


def get_data_loader(train_dir, batch_size):
    # Load MNIST dataset
    transform = transforms.Compose([
        transforms.ToTensor(),
        # transforms.Normalize((0.5,), (0.5,))
    ])

    train_dataset = torchvision.datasets.MNIST(root=train_dir, train=True, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    test_dataset = torchvision.datasets.MNIST(root=train_dir, train=False, download=True, transform=transform)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader
